from cryptography.hazmat.primitives.asymmetric import x25519
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

# Alice 生成私钥和公钥
alice_private_key = x25519.X25519PrivateKey.generate()
alice_public_key = alice_private_key.public_key()

# Bob 生成私钥和公钥
bob_private_key = x25519.X25519PrivateKey.generate()
bob_public_key = bob_private_key.public_key()

# Alice 接收到 Bob 的公钥
received_bob_public_key = x25519.X25519PublicKey.from_public_bytes(bob_public_key.public_bytes(
    encoding=Encoding.Raw,
    format=PublicFormat.Raw
))

# Bob 接收到 Alice 的公钥
received_alice_public_key = x25519.X25519PublicKey.from_public_bytes(alice_public_key.public_bytes(
    encoding=Encoding.Raw,
    format=PublicFormat.Raw
))

# Alice 计算共享密钥
alice_shared_secret = alice_private_key.exchange(received_bob_public_key)

# Bob 计算共享密钥
bob_shared_secret = bob_private_key.exchange(received_alice_public_key)

# 对共享密钥进行派生，生成最终的对称密钥（可选步骤）
def derive_key(shared_secret):
    derived_key = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=None,
        info=b'handshake data',
    ).derive(shared_secret)
    return derived_key

alice_derived_key = derive_key(alice_shared_secret)
bob_derived_key = derive_key(bob_shared_secret)

# 验证 Alice 和 Bob 生成的派生密钥是否相同
if alice_derived_key == bob_derived_key:
    print("Alice 和 Bob 生成的派生密钥相同，密钥交换成功！")
    print(f"派生密钥（十六进制）: {alice_derived_key.hex()}")
else:
    print("Alice 和 Bob 生成的派生密钥不同，密钥交换失败！")
    